#! /usr/bin/env python
# $Id: ImgSrc.py 39 2009-12-04 08:21:18Z aantin0u5 $

import xml.etree.ElementTree as ET
import ClientForm as CF
from hashlib import md5
import urllib2
import re
import codecs
import os.path
import sys
from mechanize import Browser, LinkNotFoundError

INFO_URL = "http://imgsrc.ru/cli/info.php?login=%s&passwd=%s&version=1.21"
UPLOAD_FROM = """<form method=post enctype='multipart/form-data' action=http://imgsrc.ru/cli/post.php?login=%s&passwd=%s&album_id=%s>
1: <input type=file name=u1><br>
2: <input type=file name=u2><br>
3: <input type=file name=u3><br>
4: <input type=file name=u4><br>
5: <input type=file name=u5><br>
<input type=submit>
</form>
"""

USER_URL = "http://%s.imgsrc.ru"
ALBUM_URL = "http://imgsrc.ru/%s/a%s.html?lang=en"
ALBUM_URL_PASS = "http://imgsrc.ru/main/passchk.php?ad=%s&lang=en"
AID_REG = "http\://imgsrc\.ru/%s/a([^\.]+)\.html"
AID_REG2= re.compile(r"ad=([^\&]+)\&")
ALBUM_REG = r"<td><a target=_top href=.*?/%s/a(?P<aid>\d+).html[^>]*?>(?P<title>.*?)</a></td>"
ALBUM_PASS_REG = re.compile(r"<td><a target=_top href=.*?/main\/passchk\.php\?ad=(?P<aid>\d+)(?:\&[^>]*?)?>(?P<title>.*?)</a> \(password protected\)</td>",re.DOTALL)
ALBUM_TR_REG = re.compile(r"<tr (?:class=tdd )?valign=top><td><a target=_top .*?</tr>",re.DOTALL)
PIC_NAME_REG = re.compile(r"<br>'<strong>(.+?)</strong>'")
PIC_URL_FULL_REG = re.compile(r"<a href=(.+?) target=_blank><b>view full-sized</b>")
PIC_URL_REG = re.compile(r"<img [^>]*src=(.+?\.jpg)[^>]*id=bigpic[^>]*>")
PIC_NEXT_REG = re.compile(r"<a id=new_big_url href=(.+?)>")

TAG_REG = re.compile(r"<(?P<tag>\w+)(?:\s(?P<attrs>[^>]*?))?>(?P<text>.*?)</(?P=tag)>",re.DOTALL)

class BadPasswordError(Exception):
    def __str__(self):
        return "Bad password"
    def __repr(self):
        return str(self)

class DummyFile:
    def __init__(self,content):
        self.content = content
   
    def read(self,chunk=0):
        return self.content

class ListParser :
    def __init__(self,fname):
        self.fname = fname
        f = open(fname)
        lines = f.readlines()
        self.nfilestot = len(lines) 
        lines = [l.split('::') for l in lines if l.strip()!=""]
        lines = [(l[0].strip(),l[1].strip()) for l in lines]
        f.close()
        base,ext = os.path.splitext(fname)
        trans = base+'_trans'+ext
        self.ftrans = trans
        try:
            ftrans = open(trans)
            tlines = ftrans.readlines()
            ftrans.close()
            tlines = [l.split('::') for l in tlines]
            self.tlist = dict( [ l[0].strip(),l[1].strip() ] for l in tlines )
        except IOError :
            self.tlist = {}
        self.flist = [l for l in lines if l[0] not in self.tlist.keys()]
        print "Remaining %i/%i"%(self.nfilestot-len(self.tlist),self.nfilestot)

    def getNext5(self):
        n = 0
        l = []
        while n<5 :
            try :
                f = self.flist[0]
            except IndexError :
                print "No more files..."
                break
            self.flist = self.flist[1:]
            try :
                try :
                    r = urllib2.urlopen(f[1])
                except ValueError :
                    r = open(f[1])
                cont = r.read()
            except urllib2.HTTPError :
                cont='<'
            if cont[0]=='<' :
                self.appendToTransFile([f[0]],0)
                sys.stdout.write("-")
                sys.stdout.flush()
                continue
            sys.stdout.write("+")
            df = DummyFile(cont)
            sys.stdout.flush()
            l.append((f[0],df))
            n+=1
        print
        return l
    
    def appendToTransFile(self,fnames,status):
        t = open(self.ftrans,'a')
        for fname in fnames :
            t.write("%s :: %i\n"%(fname,status))
            self.tlist[fname]=status
        t.close()
        
    def upload(self,login,password,album_id):
        user = ImgSrcPerso(login,password)
        print "Retrieving album..."
        album = user.getAlbum(album_id)
        print "Album \"%s\" retrieved !"%album.name
        while True :
            print "Fetching next files..."
            next = self.getNext5()
            if len(next)<=0 : break
            album.uploadFiles(next)
            self.appendToTransFile([f[0] for f in next],1)
            print "Remaining %i/%i"%(self.nfilestot-len(self.tlist),self.nfilestot)



class AlbumPerso :
    def __init__(self,user,element):
        self.user = user
        self.id = element.attrib["id"]
        self.name = element.find("name").text
        self.modified = element.find("modified").text
        self.photos = int(element.find("photos").text)
        try :
            self.password = element.find("password").text
        except AttributeError :
            self.password = None

    def __str__(self):
        return "AlbumPerso(id : %s, name : %s, modidied : %s, photos : %i)"%(self.id,self.name,self.modified,self.photos)
    
    def uploadFiles(self,files):
        uf = UploadForm(self.user.login, self.user.passhash,self.id)
        uf.upload(files)

class UploadForm :
    def __init__(self,login,passhash,aid):
        self.stringform = UPLOAD_FROM%(login,passhash,aid)

    def upload(self,files):
        """ Upload up to 5 files
        """
        print "Starting uploading..."
        nfiles = len(files)
        if nfiles > 5 : raise Exception("Too many files : %i"%nfiles)
         
        df = DummyFile(self.stringform)
        form = CF.ParseFile(df,"",backwards_compat=False)[0]
        file_controls = form.controls[:5]

        n = 0
        for f in files :
            c = file_controls[n]
            c.add_file(f[1],"image/jpeg",f[0])
            n+=1
        # print form
        r = form.click()
        opener = urllib2.build_opener()
        resp = opener.open(r)
        # print cont
        xml = ET.parse(resp)
        status = xml.find("status").text
        if status != 'OK' : raise Error("Problem during uploading")
        print "%i files uploaded !"%nfiles

       
class ImgSrcPerso :
    def __init__(self,login,passwd):
        self.login = login
        self.passwd = passwd
        self.passhash = md5(passwd).hexdigest()
    
    def getAlbums(self):
        url = urllib2.urlopen(INFO_URL%(self.login,self.passhash))
        et = ET.parse(url)
        albums = et.find("albums").getchildren()
        aa = [AlbumPerso(self,e) for e in albums]
        return aa
    
    def getAlbum(self,aid):
        albums = self.getAlbums()
        for a in albums :
            if a.id == aid : return a
        return None
        
    def createAlbum(self,title,passwd=None,category_id=None):
        url = INFO_URL%(self.login,self.passhash)+"&create=%s"%title
        if not passwd is None : url+="&create_passwd=%s"%passwd
        if not category_id is None : url+="&create_category=%s"%category_id
        #url = urllib2.urlopen(url)
        #print url.read()
    
def clean_url(url):
    start = "http%3A%2F%2F"
    i = url.strip().rfind(start)
    if i>0 :  url = "http://"+url[i+len(start):]
    start = "http://"
    i = url.strip().rfind(start)
    if i>0 :  url = "http://"+url[i+len(start):]
    if not url.startswith("http://"):
        url = "http://imgsrc.ru"+url
    return url
    

def album_from_tr(user,tr):
    td_reg = re.compile(r"<td[^>]*?>.+?</td>")
    td_content_reg = re.compile(r"<td[^>]*?>(?P<content>.*?)</td>")
    tds = td_reg.findall(tr)
    ALB_REG = re.compile(ALBUM_REG%user)
    match = ALB_REG.match(tds[0])
    if match is None :
        match = ALBUM_PASS_REG.match(tds[0])
        hasPass = True
    else :
        hasPass = False

    aid = match.groupdict()['aid']
    title = match.groupdict()['title']
    photos = int(td_content_reg.match(tds[1]).groupdict()['content'])
    cat = td_content_reg.match(tds[2]).groupdict()['content']
    comments = td_content_reg.match(tds[4]).groupdict()['content']
    match = TAG_REG.search(comments)
    if match : comments = int(match.groupdict()['text'])
    else : comments = 0
    modified = td_content_reg.match(tds[5]).groupdict()['content']
    match = TAG_REG.search(modified)
    if match is not None : modified = match.groupdict()['text']
    return Album(aid,title,ALBUM_URL%(user,aid),photos,cat,comments,modified,hasPass)

def parse_pict_page(html):
    pict_name = PIC_NAME_REG.findall(html)[0]
    try :
        pict_url =  PIC_URL_FULL_REG.findall(html)[0]
    except IndexError :
        pict_url =  PIC_URL_REG.findall(html)[0]
    print pict_url
    next =  PIC_NEXT_REG.findall(html)[0]
    if not next.endswith('#bp') : next = None
    return (pict_name,pict_url,next)

class Album :
    def __init__(self,aid,title,url,photos,cat,comments,modified,hasPass):
        self.title = title
        self.id = aid
        self.url = clean_url(url)
        self.photos = photos
        self.category = cat
        self.comments = comments
        self.modified = modified
        self.hasPass = hasPass
        
    def __str__(self):
        return codecs.encode("Album(id = %s, title=%s, photos=%i, category=%s, comments=%i, modified=%s, hasPass=%s)"%(self.id,self.title,self.photos,self.category,self.comments,self.modified,str(self.hasPass)),'utf8')
        
    def __repr__(self):
        return str(self)
    
    def getPhotos(self,password = ""):
        if self.hasPass and password == "" : raise BadPasswordError()
        photos = []
        b = Browser()
        
        resp = b.open(self.url)
#        print self.url
#        print resp.get_data()
        try:
            l = b.find_link(text="Continue to album >>")
            resp = b.open(clean_url(l.url))
        except LinkNotFoundError : pass
        except Exception,e:
            raise e
           
        while True :
            html = resp.get_data()
            if  html.find("has protected his work") != -1 :
                ff = b.forms()
                f = ff.next()
                f["pwd"]=password
                b.form = f
                resp = b.submit()
                #print resp.get_data()
                html = resp.get_data()
                if  html.find("has protected his work") != -1 : raise BadPasswordError()
            name,url,next = parse_pict_page(html)	    
            photos.append((name,url))
            if next is None : break
            resp = b.open(next)
        return photos


class ImgSrc :
    def __init__(self,user):
        self.user = user

    def getAlbums(self):
        url = urllib2.urlopen(USER_URL%self.user)
        html = url.read()
        trs = ALBUM_TR_REG.findall(html)
        return [album_from_tr(self.user,tr) for tr in trs]
        

def upload_list(flist,login,password,album_id):
    """ upload_list(list_file, login, password, album_id)
        
        Upload the urls/files listed in 'list_file'.
        'list_file' must have the following format :
        <name> :: <url>
        
        <url> can be either an internet url or a location 
        on a local drive.
        
        <name> is given because on some sites the name of the 
        original file is not seen in url (e.g. with some content
        managment systems).    
    """
    lp = ListParser(flist)
    lp.upload(login,password,album_id)

def display_my_albums(login,password):
    i = ImgSrcPerso(login,password)
    als = i.getAlbums()
    for a in als :
        print a

def display_user_albums(user):
    i=ImgSrc(user)
    als=i.getAlbums()
    print len(als),"album(s) found."
    for a in als :
        print a

def get_user_albums(user):
    i=ImgSrc(user)
    als=i.getAlbums()
    return dict([ (a.id, a) for a in als ])

def get_user_album(user,aid):
    als = get_user_albums(user)
    return als[aid]

def get_photos_from_album(user,aid,password=""):
    print "Opening album %s"%aid
    album = get_user_album(user,aid)
    print "Album %s found with title '%s'"%(album.id,album.title)
    print "%i photos in the album"%album.photos
    print "Now listing album content..."
    return album.getPhotos(password)
    

def write_photos_list(filename,user,aid,password=""):
    """ write_photos_list(filename,user,aid,password="")
        
        Writes a file containing the list of the photos
        from the given album into a format readable by
        'upload_list'
    """
    photos = get_photos_from_album(user,aid,password)
    f = open(filename,'w')
    for name,url in photos :
        f.write("%s :: %s\n"%(name,url))
    f.close()
def usage():
    print """
    usage : python Imgsrc.py <command> <arguments>
    
    available commands :
        - albums <user> : 
            displays the albums description for the user <user>.
            
        - list <filename> <user> <album_id> [<password>]
            
        - upload <login> <password> <album_id> <list_file> :
            Upload the files in <list_file> to the album with id
            <album_id>. Since it is your account you need to give
            your <login> and <password> information.
            <list_file> contains the list of urls to download.
            There must be one line per image and each line has the following
            format :
                <name> :: <url>
            <name> is needed because the file name might not correspond to
            the one given in url. This can be the case when the server uses
            a content system manager. This is the case for imgsrc.ru site
            for instance.
        """
    sys.exit(-1)

def main():
    try :
        command = sys.argv[1]
    except IndexError :
        usage()
    args = sys.argv[2:]
    try:
        if command == 'albums':
            user = args[0]
            print "Albums for user %s"%user
            display_user_albums(user)
          
        elif command == 'list' :
            try :
                filename,user,album_id,password = args
            except ValueError :
                filename,user,album_id = args
                password = ""
            write_photos_list(filename,user,album_id,password)
       
        elif command == 'upload' :
            login,password,album_id,list_file = args
            upload_list(list_file,login,password,album_id)
            
        else : usage()
    except ValueError : usage()         

if __name__ == '__main__' :
    main()
